import argparse
import torch
import gendoc as gendoc
from utils.conf import GLOBAL_CONF, read_conf, update_conf
from utils.consts import LOG_PATH
from utils.tools import get_timestamp, write_json, check_path
from trainer.stageAB import Atrainer, Btrainer
from trainer.stageG import Gtrainer


DEBUG = {
    'debug': True,
    'conf': 'gca',
    'gpu': 1
}


def release():
    # from cmdline get config.ini's name
    parser = argparse.ArgumentParser(description='data and trianing configs')
    parser.add_argument('--ini', '-i', type=str, required=True)
    parser.add_argument('--gpu', '-g', type=int, required=True)
    args = parser.parse_args()

    # incase cpu overload
    torch.set_num_threads(1)

    # read config.ini
    flag = read_conf(args.ini)
    if not flag:
        print("without setting task or set invalid task name")
        return
    # set_gpu_id
    GLOBAL_CONF['gpu'] = args.gpu


def debug():
    flag = read_conf(DEBUG['conf'])
    if not flag:
        print("without setting task or set invalid task name")
        return
    # set_gpu_id
    GLOBAL_CONF['gpu'] = DEBUG['gpu']


def main():
    if DEBUG['debug']:
        debug()
    else:
        release()

    # generate taskid
    taskid = '-'.join([GLOBAL_CONF['task'][-1], GLOBAL_CONF['dataset'],
                      GLOBAL_CONF['distr_flag'], str(get_timestamp())])

    # log task configer settings
    check_path(LOG_PATH)

    trainer = None
    # according task flag to start a train task
    if GLOBAL_CONF['task'] == 'stageA':
        trainer = Atrainer(task_id=taskid, conf=GLOBAL_CONF, title_icon='🔥')
    elif GLOBAL_CONF['task'] == 'stageB':
        trainer = Btrainer(task_id=taskid, conf=GLOBAL_CONF, title_icon='🌿')
    elif GLOBAL_CONF['task'] == 'stageG':
        trainer = Gtrainer(task_id=taskid, conf=GLOBAL_CONF, title_icon='🧊')
    else:
        return
    res_dict = trainer.train()

    # use res_dict to update globaldict then log all info
    update_conf(res_dict)
    write_json(f"{LOG_PATH}/{taskid}.json", GLOBAL_CONF)

    # draw and save imgs
    gendoc.draw(taskid, GLOBAL_CONF)
    # gen docs
    gendoc.gen_doc(taskid, GLOBAL_CONF)
    print(f"\n===> Exp Doc: 'docs/{taskid}.md'")
    # log task id
    with open("docs/tasklists.log", 'a') as f:
        f.write(taskid+'\n')


if __name__ == '__main__':
    import warnings
    warnings.filterwarnings("ignore", category=DeprecationWarning)
    main()
